package com.sofun.ai.controller;

import com.sofun.ai.assistant.RagAssistant;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@RestController()
@RequestMapping("/embedding")
public class EmbeddingController {

     @Resource
     QwenEmbeddingModel qwenEmbeddingModel;

    @Resource
    QwenChatModel qwenChatModel;

/*    @Resource
    EmbeddingStore<TextSegment> embeddingStore;*/

     @Resource
     RagAssistant ragAssistant;

     @RequestMapping(value = "/embed")
     public String embedding(@RequestParam(defaultValue="我叫张佳伟") String message) {
         // 文本向量化
         Response<Embedding> embed = qwenEmbeddingModel.embed(message);
         return embed.toString();
     }

    @RequestMapping(value = "/embed/search")
    public void embeddingSearch()  {

        // 适应InMemory向量数据库，embedding阶段
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        // 利用向量模型进行向量化， 然后存储向量到向量数据库
        TextSegment orderSegment = TextSegment.from("""
                预订航班:
                - 通过我们的网站或移动应用程序预订。
                - 预订时需要全额付款。
                - 确保个人信息（姓名、ID 等）的准确性，因为更正可能会产生 25 的费用。
                """);
        Embedding orderEmbedding = qwenEmbeddingModel.embed(orderSegment).content();
        // 存入向量数据库
        embeddingStore.add(orderEmbedding, orderSegment);

        // 利用向量模型进行向量化， 然后存储向量到向量数据库
        TextSegment cancelSegment = TextSegment.from("""
                取消预订:
                - 最晚在航班起飞前 48 小时取消。
                - 取消费用：经济舱 75 美元，豪华经济舱 50 美元，商务舱 25 美元。
                - 退款将在 7 个工作日内处理。
                """);
        Embedding cancelEmbedding = qwenEmbeddingModel.embed(cancelSegment).content();
        embeddingStore.add(cancelEmbedding, cancelSegment);


        // ----------------------数据检索阶段------------------------------

        // 需要查询的内容 向量化
        Embedding queryEmbedding = qwenEmbeddingModel.embed("退票要多少钱").content();

        // 去向量数据库查询
        // 构建查询条件
        EmbeddingSearchRequest build = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                //相似性检索，只有满足相似度score的才会检索出来(相似度最大为1)
//                .maxResults(1)
//                .minScore(0.8)
                .build();

        // 查询
        EmbeddingSearchResult<TextSegment> segmentEmbeddingSearchResult = embeddingStore.search(build);

        // 按score排序并获取score最大的结果
        // 降序排序
        segmentEmbeddingSearchResult.matches().stream().min((match1, match2) -> Double.compare(match2.score(), match1.score())) // 获取第一个元素（score最大的）
                .ifPresent(bestMatch -> {
                    System.out.println("最佳匹配结果 Score: " + bestMatch.score());
                    System.out.println("最佳匹配文本: " + bestMatch.embedded().text());
                });


        // 打印所有匹配
        System.out.println("\n所有匹配结果:");
        segmentEmbeddingSearchResult.matches().forEach(embeddingMatch -> {
            System.out.println(embeddingMatch.score()); // 0.8144288515898701
            System.out.println(embeddingMatch.embedded().text()); // I like football.
        });
    }

    @RequestMapping(value = "/embed/split")
    public void embeddingWithSplit(){
        Document document = ClassPathDocumentLoader.loadDocument("rag/terms-of-service.txt", new TextDocumentParser());



        //将文本内容分词，方便大模型理解
        DocumentByRegexSplitter splitter = new DocumentByRegexSplitter(
                "\\n\\d+\\.",  // 匹配 "1. 标题" 格式
                "\n",            // 保留换行符作为段落连接符
                100,             // 每个段最多 500 字符
                20,               // 段间重叠 50 字符以保持连贯性
                new DocumentByCharacterSplitter(100,50)
        );

        List<TextSegment> segments = splitter.split(document);

        System.out.println("分词后的数据内容:" + segments);

        // 向量化
        List<Embedding> embeddings = qwenEmbeddingModel.embedAll(segments).content();
        System.out.println(embeddings);

        //分词后的数据给向量数据库
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings,segments);

        // 生成向量检索条件
        Response<Embedding> embed = qwenEmbeddingModel.embed("退费费用");
        EmbeddingSearchRequest build = EmbeddingSearchRequest.builder()
                .queryEmbedding(embed.content())
                .maxResults(1)
                .build();
        // 查询
        EmbeddingSearchResult<TextSegment> results = embeddingStore.search(build);
        for (EmbeddingMatch<TextSegment> match : results.matches()) {
            System.out.println(match.embedded().text() + ",分数为：" + match.score());

        }
    }

    /**
     * 将向量化的数据存在向量数据库，一起交给大模型，
     * */
    @RequestMapping(value = "/embed/rag")
    public String chatWithRag(){

        Document document = ClassPathDocumentLoader.loadDocument("rag/terms-of-service.txt", new TextDocumentParser());

        //将文本内容分词，方便大模型理解
        DocumentByRegexSplitter splitter = new DocumentByRegexSplitter(
                "\\n\\d+\\.",  // 匹配 "1. 标题" 格式
                "\n",            // 保留换行符作为段落连接符
                100,             // 每个段最多 500 字符
                20,               // 段间重叠 50 字符以保持连贯性
                new DocumentByCharacterSplitter(100,50)
        );

        List<TextSegment> segments = splitter.split(document);

        System.out.println("分词后的数据内容:" + segments);

        // 向量化
        List<Embedding> embeddings = qwenEmbeddingModel.embedAll(segments).content();
        System.out.println(embeddings);

        //分词后的数据给向量数据库
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings,segments);

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(1) // 最相似的5个结果
                //.minScore(0.7) // 只找相似度在0.6以上的内容
                .build();

        RagAssistant assistant = AiServices.builder(RagAssistant.class)
                .chatLanguageModel(qwenChatModel)
                .contentRetriever(contentRetriever)
                .build();

        return assistant.chat("退费费用");
    }


    /**
     * 启动时，将向量化的数据提前加载到向量数据库，一起交给大模型，
     * */
/*    @RequestMapping(value = "/embed/rag/pre")
    public String chatWithRagPreLoad(){

        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(embeddingStore)
                .embeddingModel(qwenEmbeddingModel)
                .maxResults(1) // 最相似的5个结果
                //.minScore(0.7) // 只找相似度在0.6以上的内容
                .build();



        return ragAssistant.chat("退费费用");
    }*/
}
